
class Node:
    def __init__( self, pt, splitaxis ):
        self.point = pt
        self.splitaxis = splitaxis
        self.left = None
        self.right = None
        
def kd_dist_sqd( p1, p2 ):
    d = 0
    for pair in zip( p1,p2 ):
        x = pair[0] - pair[1]
        d += x * x
    return d

def kd_insert( tree, point ):
    
    dimen = len(point)
    
    axis = tree.splitaxis
    is_left = tree.point[axis] > point[axis]
    
    if is_left:
        if tree.left == None:
            tree.left = Node( point, (axis + 1) % dimen )
        else:
            kd_insert( tree.left, point )
    else:
        if tree.right == None:
            tree.right = Node( point, (axis + 1) % dimen )
        else:
            kd_insert( tree.right, point )

def kd_find_parent( tree, target ):

    #print( "kd_find_parent {0} {1}".format( tree.point, target ) )
    
    parent = None
    next = tree
    
    while next <> None:
        
        axis = next.splitaxis
        parent = next
        
        #print( "-kd_find_parent {0} {1}".format( next.point, axis ) )
        
        is_left = next.point[axis] > target[axis]
        
        if is_left:
            next = next.left
        else:
            next = next.right
            
    return parent
    

def kd_check_subtree( tree, target, find_state ):
    if tree == None:
        return
    
    d = kd_dist_sqd( tree.point, target )
    
    #print( "tree{0} find_state {1}, {2}; d={3}".format( tree.point, find_state[0].point, find_state[1], d ) )
    
    if ( d < find_state[1] ):
        #print( "-d < find_state[1] {0} {1}".format( tree.point, d ) )
        find_state[0] = tree
        find_state[1] = d
    
    axis = tree.splitaxis
    
    d = tree.point[axis] - target[axis]
    
    if d*d > find_state[1]:
        is_left = tree.point[axis] > target[axis]
        
        if is_left:
            kd_check_subtree( tree.left, target, find_state )
        else:
            kd_check_subtree( tree.right, target, find_state )
            
    else:
        kd_check_subtree( tree.left, target, find_state )
        kd_check_subtree( tree.right, target, find_state )
    
def kd_find_nearest( tree, target ):

    p = kd_find_parent( tree, target )
    d = kd_dist_sqd( p.point, target )
    
    find_state = [p, d]
    
    kd_check_subtree( tree, target, find_state )
    
    return find_state
    
def kd_print( tree, indent="", type="" ):
    if tree == None:
        return
        
    print( "{0}{1}{2} split at {3}".format( indent, type, tree.point, tree.splitaxis ) )    
    kd_print( tree.left, indent + " ", "l" )
    kd_print( tree.right, indent + " ", "r" )
    
if __name__ == "__main__":
    
    import random
    
    random.seed(12345)    
    points = []
    
    for i in range(8):
        points.append( (random.randint(1,100), random.randint(1,100)) )
        
    print points
    
    #target = (points[2][0]-1, points[2][1])
    target = (random.randint(1,100), random.randint(1,100))
    
    dimen = len(points[0])
    print( "dimension={0}".format( dimen ) )
    
    tree = Node(points[0], 0)
    
    for i in range(1, 8):
        kd_insert( tree, points[i] )
        print( points[i], kd_dist_sqd( target, points[i] ) )
    
    #kd_print( tree )
    
    n = kd_find_nearest( tree, target )
    
    print( target, n[0].point, n[1] )
    #print kd_dist_sqd( points[0], points[1] )
    #print kd_dist_sqd( points[1], points[1] )
    #print (tree.left.point if tree.left <> None else None)
    #print (tree.right.point if tree.right <> None else None)